"""
    2/20
    Since normal map given by RealFaceDB is in object-space, we should convert then into tangent-space via shape texture
    stored in ReadFaceDB which indicate the coordinate of shape in object-space that transform into [0, 255].

    Get the value from shape-patch's corner, and calculate the TBN matrix.
    Considering that the shape represented by this shape-patch map may be curved, the result may not accurate.

    available variant:
    scalar_rgb, scalar_spectral, cuda_ad_rgb, llvm_ad_rgb
"""

import time

import mitsuba as mi
import numpy as np

import config

mi.set_variant('cuda_ad_rgb')


def find_tbn(shape_filepath) -> np.ndarray:
    """
        Calculate TBN of the given shape-patch.

        :param shape_filepath : the given shape-patch filepath.
    """
    '''
        Set variant
    '''
    mi.set_variant('cuda_ad_rgb')

    '''
        Obtain base information
    '''
    bmp = mi.Bitmap(shape_filepath)             # load shape-patch from file
    bmp_np = np.array(bmp)                      # convert into NpArray
    height, width, channel = bmp_np.shape       # get the shape of patch

    '''
        Get three sample points from shape-map.
        Remember to convert data range of [0, 255] into [0, 1] before passing to following process.
    '''
    origin = bmp_np[height - 1, 0] / config.PNG_CHANNELS
    vertex_b = bmp_np[0, 0] / config.PNG_CHANNELS
    vertex_t = bmp_np[height - 1][width - 1] / config.PNG_CHANNELS

    '''
        Calculate basis of TBN.
    '''
    tangent = vertex_t - origin
    bi_tangent = vertex_b - origin
    normal = np.cross(tangent, bi_tangent)
    bi_tangent = np.cross(normal, tangent)      # tangent and bi_tangent may not orthogonal, so recompute one of them.

    '''
        Calculate TBN.
        Don't forget to normalize basis of TBN before calculate TBN.
        * Using [np.linalg.norm] to compute length of vector.
    '''
    tangent = tangent / np.linalg.norm(tangent, 2)
    bi_tangent = bi_tangent / np.linalg.norm(bi_tangent, 2)
    normal = normal / np.linalg.norm(normal, 2)
    tbn = np.array([tangent, bi_tangent, normal])

    return tbn


def convert_normal_map(normal_map_filepath, shape_map_filepath) -> mi.TensorXf:
    """
        Given a normal map's filepath and corresponding shape-patch map's filepath,
        we will read the content of file and return tangent space normal map.

        :return : result tangent map in TensorXf
    """
    '''
        Set variant
    '''
    mi.set_variant('cuda_ad_rgb')

    '''
        Obtain necessary information for normal map convert.
    '''
    normal_bmp = mi.Bitmap(normal_map_filepath)
    tbn = find_tbn(shape_map_filepath)
    normal_np = np.array(normal_bmp)

    '''
        The normal map loading from PNG file has been converted into [0, 255], however their real data range should be 
        [-1 , 1]. A convert need to be applied to loaded normal map before using it.
    '''
    normal_np = (normal_np / config.PNG_CHANNELS) * 2 - 1

    '''
        Convert normal into tangent-space via TBN and Einstein-Sum.
    '''
    tangent_normal_np = np.einsum('kij, lj->kil', normal_np, tbn)

    '''
        Normalize the normal map via [np.linalg.norm]
    '''
    length_matrix = np.linalg.norm(tangent_normal_np, 2, axis=-1)       # Valid vector is in the last axis of ndarray.
    # The length-matrix shape is (width, height), containing the length of each vector in normal map. A dim expand
    # should be applied to make the element-wise operate correctly. Note that the axis '-1' indicate the last axis of
    # matrix, which the valid vector located in.
    normalized_tangent_normal_map = tangent_normal_np / np.expand_dims(length_matrix, -1)

    tangent_normal_tensor = mi.TensorXf(normalized_tangent_normal_map)
    if config.DEBUG:
        # print("Normal map before convert:")
        # print(normal_np)
        # print("Normalized tangent normal map has been shown below:")
        # print(normalized_tangent_normal_map)
        '''
            DEBUG : Save result into a temporary file and return the temp filepath.
            In order to use in OpenGL, the map would need to convert to [0, 255] format.
        '''
        temp_filepath = config.temporary_filepath
        identity = int(time.time())
        temp_filepath = temp_filepath + str(identity - 1676947965) + '.exr'
        tangent_normal_bmp = mi.Bitmap(tangent_normal_tensor)
        tangent_normal_bmp.write(mi.filesystem.path(temp_filepath))

    return tangent_normal_tensor


if __name__ == "__main__":
    pass
